from ..arch import Arch, from_rsf
from typing import List, Tuple
import sys

# mto: 一个 arch 转换为另一个 arch 的最小操作次数
# source target 的 cluster 数量相等：不需要增加/删除 cluster
# 对于每个 cluster 中没有匹配的 obj，如果 target 中不存在，则删去（2次操作）
# 如果 target 中存在，则移动至正确的 cluster（1次操作）
# 剩余的再新增（2次操作）
# 如果 source 的 cluster 数量更多，则先将未匹配的 cluster 内容移出然后删除 cluster
# 如果 source 的 cluster 数量更少，则增加 cluster，再移动内容


class KM:
    def __init__(self, arch1: Arch, arch2: Arch):
        self.n = max(len(arch1), len(arch2))

        self.key2vertex_1 = dict()
        self.vertex2key_1 = dict()
        self.key2vertex_2 = dict()
        self.vertex2key_2 = dict()

        for i, k in enumerate(sorted(arch1.keys())):
            self.key2vertex_1[k] = i
            self.vertex2key_1[i] = k
        for i, k in enumerate(sorted(arch2.keys())):
            self.key2vertex_2[k] = i
            self.vertex2key_2[i] = k

        self.weight = []
        self.lx = [0] * self.n
        self.ly = [0] * self.n

        for i in range(self.n):
            temp = [0] * self.n
            self.weight.append(temp)

        for i, k1 in self.vertex2key_1.items():
            for j, k2 in self.vertex2key_2.items():
                self.weight[i][j] = len(arch1[k1] & arch2[k2])

        for i in range(self.n):
            self.lx[i] = max(self.weight[i])

        self.match = [-1] * self.n
        self.S = [False] * self.n
        self.T = [False] * self.n

    def update(self):
        a = 2 ** 31
        for i in [x for x in range(self.n) if self.S[x]]:
            for j in [x for x in range(self.n) if not self.T[x]]:
                a = min(a, self.lx[i] + self.ly[j] - self.weight[i][j])
        for i in [x for x in range(self.n) if self.S[x]]:
            self.lx[i] -= a
        for i in [x for x in range(self.n) if self.T[x]]:
            self.ly[i] += a

    def find(self, i: int) -> bool:
        self.S[i] = True
        for j in [x for x in range(self.n) if not self.T[x]]:
            if self.lx[i] + self.ly[j] == self.weight[i][j]:
                self.T[j] = True
                if self.match[j] == -1 or self.find(self.match[j]):
                    self.match[j] = i
                    return True
        return False

    def km(self):
        for i in range(self.n):
            while True:
                self.S = [False] * self.n
                self.T = [False] * self.n
                if not self.find(i):
                    self.update()
                else:
                    break

    def get_result(self) -> List[Tuple[str, str, int]]:
        res = []
        for j, i in enumerate(self.match):
            try:
                source = self.vertex2key_1[i]
            except KeyError:
                source = "Null"
            try:
                target = self.vertex2key_2[j]
            except KeyError:
                target = "Null"
            res.append((source, target, self.weight[i][j]))
        return res

    def get_match_num(self) -> int:
        res = 0
        for j, i in enumerate(self.match):
            res += self.weight[i][j]
        return res


def mto(arch1: Arch, arch2: Arch):
    res = 0
    km = KM(arch1, arch2)
    km.km()
    res += abs(len(arch1) - len(arch2))

    objs1 = set()
    for _, cluster in arch1.items():
        objs1.update(cluster)
    objs2 = set()
    for _, cluster in arch2.items():
        objs2.update(cluster)

    # 出现在 1 但没有出现在 2
    a = len(objs1 - objs2)
    # 1 2 共有的部分
    b = len(objs1 & objs2)
    # 出现在 2 但没有出现在 1
    c = len(objs2 - objs1)

    assert b >= km.get_match_num()

    res += a * 2 + (b - km.get_match_num()) + c * 2
    return res


def aco(arch: Arch) -> int:
    res = 0
    res += len(arch)
    objs = set()
    for _, cluster in arch.items():
        objs.update(cluster)
    res += len(objs) * 2
    return res


def a2a(arch1: Arch, arch2: Arch) -> float:
    sys.setrecursionlimit(1000000)
    res = mto(arch1, arch2) / (aco(arch1) + aco(arch2))
    return 1 - res


if __name__ == '__main__':
    arch1 = from_rsf("../bash_compatible_with_gt.rsf")
    arch2 = from_rsf("../bash_gt.rsf")
    print(a2a(arch1, arch2))
